package edu.uaskl.jqf.algorithm;

import java.util.Arrays;

import edu.uaskl.jqf.algorithm.interfaces.Function;
import edu.uaskl.jqf.model.Complex;
import edu.uaskl.jqf.model.Quregister;
import edu.uaskl.jqf.model.matrix.MatrixComplex;
import edu.uaskl.jqf.tools.MathTools;

public class QftModularExponentation implements Function {

    private final int base;
    private final int modulus;
    private final int numberOfQftBits;
    private final int numberOfModBits;

    /**
     * Prepares the modular exponentiation, x^a mod n<br>
     * where x = base and n = modulus
     * 
     * @param base
     *            is the x in x^a mod n
     * @param modulus
     *            is the n in x^a mod n
     * @param numberOfQftBits
     *            the upper bits to do the qft mod exp
     * @param numberOfModBits
     *            the lower bits to do the qft mod exp
     */
    public QftModularExponentation(final int base, final int modulus, final int numberOfQftBits, final int numberOfModBits) {
        this.base = base;
        this.modulus = modulus;
        this.numberOfQftBits = numberOfQftBits;
        this.numberOfModBits = numberOfModBits;
    }

    @Override
    public Quregister apply(final Quregister quRegister) {
        return apply(quRegister.getQubitRegister());
    }

    public Quregister apply(final MatrixComplex qubitRegister) {
        checkDimensions(qubitRegister);
        System.out.println("base: " + base + ", modulus: " + modulus + ", dimension: " + qubitRegister.getDimensionAsString());

        final Complex[][] resultArray = new Complex[qubitRegister.getRows()][qubitRegister.getColumns()]; // TODO columns must be 1? change to transposed array?
        for (final Complex[] resultArrayItem : resultArray)
            Arrays.fill(resultArrayItem, Complex.get00());

        // final long start = System.currentTimeMillis();
        for (int i = 0; i < qubitRegister.getRows(); i++) {
            final Complex currentElement = qubitRegister.get(i, 0);
            if (currentElement.equals(Complex.get00()))
                continue;

            final int upperPart = i >> numberOfModBits;
            final int modExpValue = MathTools.modularExponentiationBaseExpMod(base, upperPart, modulus);

            final int newIndex = (upperPart << numberOfModBits) + modExpValue;
            resultArray[newIndex][0] = currentElement;
        }
        // final long end = System.currentTimeMillis();
        // System.out.println("duration: " + (end - start) + "ms");

        return new Quregister(MatrixComplex.createMatrixBasedOnArray2d(resultArray));
    }

    private void checkDimensions(final MatrixComplex qubitRegister) {
        if (qubitRegister.getRows() != (1 << (numberOfQftBits + numberOfModBits)))
            throw new IllegalStateException("Dimensions mismatch. The given number of qubits does not correlate to the size of the qubitregister.\n"
                    + "numberOfQftBits: " + numberOfQftBits + ", numberOfModBits: " + numberOfModBits + ", sum: " + (numberOfQftBits + numberOfModBits)
                    + ", numbers up to: " + (1 << (numberOfQftBits + numberOfModBits)) + "\n" + "Rows from qubitRegister: " + qubitRegister.getRows());
    }
}
